package com.sojson.util.search.es.util;

import java.util.ArrayList;
import java.util.List;
import java.util.Map;
import java.util.Map.Entry;
import java.util.Objects;
import java.util.Set;

import org.elasticsearch.action.search.SearchRequest;
import org.elasticsearch.action.search.SearchResponse;
import org.elasticsearch.client.RequestOptions;
import org.elasticsearch.client.RestHighLevelClient;
import org.elasticsearch.index.query.AbstractQueryBuilder;
import org.elasticsearch.index.query.BoolQueryBuilder;
import org.elasticsearch.index.query.QueryBuilders;
import org.elasticsearch.script.Script;
import org.elasticsearch.search.aggregations.AggregationBuilder;
import org.elasticsearch.search.builder.SearchSourceBuilder;
import org.elasticsearch.search.sort.FieldSortBuilder;
import org.elasticsearch.search.sort.SortBuilder;
import org.elasticsearch.search.sort.SortBuilders;
import org.elasticsearch.search.sort.SortOrder;
import org.springframework.data.domain.PageRequest;
import org.springframework.data.elasticsearch.core.ElasticsearchOperations;
import org.springframework.data.elasticsearch.core.IndexOperations;
import org.springframework.data.elasticsearch.core.SearchHit;
import org.springframework.data.elasticsearch.core.SearchHits;
import org.springframework.data.elasticsearch.core.document.Document;
import org.springframework.data.elasticsearch.core.query.NativeSearchQuery;
import org.springframework.data.elasticsearch.core.query.ScriptField;
import org.springframework.data.elasticsearch.repository.ElasticsearchRepository;

import com.github.pagehelper.Page;
import com.sojson.base.impl.BasePageDto;
import com.sojson.base.impl.OrderItem;
import com.sojson.util.SpringUtil;
import com.sojson.util.StringUtil;
import com.sojson.util.annotation.AnnoUtil;
import com.sojson.util.page.PageUtil;
import com.sojson.util.search.es.bean.EsBaseBean;

import cn.hutool.core.collection.CollUtil;
import lombok.extern.slf4j.Slf4j;

/**
 * ES工具类
 * 
 * @author liu
 * @date 2022-07-25
 */
@Slf4j
public class EsUtil {

    private static final ElasticsearchOperations ES = SpringUtil.getBeanByType(ElasticsearchOperations.class);
    private static final RestHighLevelClient CLIENT = SpringUtil.getBeanByType(RestHighLevelClient.class);
    /** 索引名称在注解中的前缀 */
    public static final String INDEX_NAME_ANNO_PREFIX = AnnoUtil.VALUE_CLASS_PRIFIX
        //
        + "com.sojson.util.search.es.util.EsUtil"
        //
        + AnnoUtil.VALUE_METHOD_PRIFIX
        //
        + "getIndexName"
        //
        + AnnoUtil.VALUE_PARAM_PRIFIX;

    /**
     * 创建索引
     * 
     * @param clazz
     * @return
     */
    public static boolean createIndex(Class<?> clazz) {
        return createIndex(clazz, false);
    }

    /**
     * 创建索引
     * 
     * @param clazz
     * @param isReInit  是否重新初始化标识
     * @return
     */
    public static boolean createIndex(Class<?> clazz, boolean isReInit) {
        IndexOperations index = ES.indexOps(clazz);
        // 是否创建成功标识
        boolean create = false;

        // 1. 判断索引是否存在
        boolean exists = index.exists();
        // 2. 若需要重新初始化并且索引存在,则删除索引,删除索引成功后,将索引存在标识设置为false
        if (isReInit && exists && index.delete()) {
            // 删除索引
            exists = false;
            log.info("ES索引: " + clazz.getName() + "删除成功!");
        }
        // 3.若索引不存在则创建索引
        if (!exists) {
            create = index.create();
        }
        // 4.创建成功后,设置mapping
        if (create) {
            log.info("ES索引: " + clazz.getName() + "创建成功!");
            Document mapping = index.createMapping(clazz);
            if (CollUtil.isNotEmpty(mapping)) {
                index.putMapping(mapping);
                log.info("ES索引: " + clazz.getName() + "映射成功!");
            }
        }
        return create;
    }

    /**
     * 删除索引
     * 
     * @param clazz
     * @return
     */
    public static boolean deleteIndex(Class<?> clazz) {
        IndexOperations index = ES.indexOps(clazz);
        boolean delete = index.delete();
        if (delete) {
            log.info("ES索引: " + clazz.getName() + "删除成功!");
        }
        return delete;
    }

    /**
     * 保存数据
     * 
     * @param <T>
     * @param <ID>
     * @param repository
     * @param bean
     */
    public static <T, ID> void save(ElasticsearchRepository<T, ID> repository, T bean) {
        repository.save(bean);
    }

    /**
     * 保存数据
     * 
     * @param <T>
     * @param <ID>
     * @param repository
     * @param beans
     */
    public static <T, ID> void save(ElasticsearchRepository<T, ID> repository, List<T> beans) {
        // 将数据分批插入
        List<List<T>> split = CollUtil.split(beans, 3000);
        for (List<T> list : split) {
            repository.saveAll(list);
        }
    }

    /**
     * 通过ID删除数据
     * 
     * @param <T>
     * @param <ID>
     * @param repository
     * @param id
     */
    public static <T, ID> void deleteById(ElasticsearchRepository<T, ID> repository, ID id) {
        repository.deleteById(id);
    }

    /**
     * 通过ID删除数据
     * 
     * @param <T>
     * @param <ID>
     * @param repository
     * @param bean
     */
    public static <T, ID> void delete(ElasticsearchRepository<T, ID> repository, T bean) {
        repository.delete(bean);
    }

    /**
     * 通过ID删除数据
     * 
     * @param <T>
     * @param <ID>
     * @param repository
     * @param beans
     */
    public static <T, ID> void delete(ElasticsearchRepository<T, ID> repository, List<T> beans) {
        repository.deleteAll(beans);
    }

    /**
     * 清空数据
     * 
     * @param <T>
     * @param <ID>
     * @param repository
     */
    public static <T, ID> void delete(ElasticsearchRepository<T, ID> repository) {
        repository.deleteAll();
    }

    /**
     * 通过ID判断数据是否存在
     * 
     * @param <T>
     * @param <ID>
     * @param repository
     * @param id
     * @return
     */
    public static <T, ID> boolean exists(ElasticsearchRepository<T, ID> repository, ID id) {
        return repository.existsById(id);
    }

    /**
     * 查询数据
     * 
     * @param <T>
     * @param clazz     索引类
     * @return
     * @throws Exception
     */
    public static <T> List<T> search(Class<T> clazz) throws Exception {
        return search(clazz, new BasePageDto());
    }

    /**
     * 查询数据
     * 
     * @param <T>
     * @param clazz     索引类
     * @param builder   查询条件{@link com.sojson.util.search.es.util.EsSearchUtil#build()}
     * @return
     * @throws Exception
     */
    public static <T> List<T> search(Class<T> clazz, AbstractQueryBuilder<?> builder) throws Exception {
        return search(clazz, builder, new BasePageDto());
    }

    /**
     * 查询数据
     * 
     * @param <T>
     * @param clazz     索引类
     * @param page      分页和排序
     * @return
     * @throws Exception
     */
    public static <T> List<T> search(Class<T> clazz, BasePageDto page) throws Exception {
        return search(clazz, QueryBuilders.boolQuery(), page);
    }

    /**
     * 查询数据
     * 
     * @param <T>
     * @param clazz     索引类
     * @param scripts   脚本字段(key为返回的字段名)(由其他字段计算得到){@link com.sojson.util.search.es.util.EsSearchUtil#createScript()}
     * @return
     * @throws Exception
     */
    public static <T> List<T> search(Class<T> clazz, Map<String, Script> scripts) throws Exception {
        return search(clazz, QueryBuilders.boolQuery(), scripts);
    }

    /**
     * 查询数据
     * 
     * @param <T>
     * @param clazz     索引类
     * @param builder   查询条件{@link com.sojson.util.search.es.util.EsSearchUtil#build()}
     * @param page      分页和排序
     * @return
     * @throws Exception
     */
    public static <T> List<T> search(Class<T> clazz, AbstractQueryBuilder<?> builder, BasePageDto page)
        throws Exception {
        return search(clazz, builder, page, null, false);
    }

    /**
     * 查询数据
     * 
     * @param <T>
     * @param clazz     索引类
     * @param builder   查询条件{@link com.sojson.util.search.es.util.EsSearchUtil#build()}
     * @param scripts   脚本字段(key为返回的字段名)(由其他字段计算得到){@link com.sojson.util.search.es.util.EsSearchUtil#createScript()}
     * @return
     * @throws Exception
     */
    public static <T> List<T> search(Class<T> clazz, AbstractQueryBuilder<?> builder, Map<String, Script> scripts)
        throws Exception {
        return search(clazz, builder, new BasePageDto(), scripts, false);
    }

    /**
     * 查询数据
     * 
     * @param <T>
     * @param clazz     索引类
     * @param page      分页和排序
     * @param scripts   脚本字段(key为返回的字段名)(由其他字段计算得到){@link com.sojson.util.search.es.util.EsSearchUtil#createScript()}
     * @return
     * @throws Exception
     */
    public static <T> List<T> search(Class<T> clazz, BasePageDto page, Map<String, Script> scripts)
        throws Exception {
        return search(clazz, QueryBuilders.boolQuery(), page, scripts, false);
    }

    /**
     * 查询数据
     * 
     * @param <T>
     * @param clazz         索引类
     * @param builder       查询条件{@link com.sojson.util.search.es.util.EsSearchUtil#build()}
     * @param page          分页和排序
     * @param scripts       脚本字段(key为返回的字段名)(由其他字段计算得到){@link com.sojson.util.search.es.util.EsSearchUtil#createScript()}
     * @param isSortLoca    是否按坐标排序
     * @return
     * @throws Exception
     */
    public static <T> List<T> search(Class<T> clazz, AbstractQueryBuilder<?> builder, BasePageDto page,
        Map<String, Script> scripts, boolean isSortLoca) throws Exception {
        // 包装build
        builder = packageBuild(builder);

        // 要返回的结果
        Page<T> list = new Page<>();

        // 设置排序规则
        List<SortBuilder<?>> copySort = copySort(page.getSearchSort());

        // 设置查询条件
        NativeSearchQuery query = new NativeSearchQuery(builder, null, copySort);

        // 设置分页
        // 第几页
        Integer pageNum = page.getPageNum();
        // 显示记录数
        Integer pageSize = page.getPageSize();
        boolean isPage = pageNum != null && pageSize != null;
        if (isPage) {
            if (pageNum < 1) {
                pageNum = PageUtil.PAGE_NUM_DEFAULT;
            }
            if (pageSize < 1) {
                pageSize = PageUtil.PAGE_SIZE_DEFAULT;
            }
            PageRequest of = PageRequest.of(pageNum - 1, pageSize);
            query.setPageable(of);
            // 设置允许超过1万条查询限制
            query.setTrackTotalHits(true);
        }

        if (log.isDebugEnabled()) {
            log.debug("{{}:{}}", "\"query\"", Objects.requireNonNull(query.getQuery()));
        }

        // 设置脚本字段
        if (scripts != null && scripts.size() > 0) {
            Set<Entry<String, Script>> entrySet = scripts.entrySet();
            for (Entry<String, Script> entry : entrySet) {
                ScriptField scriptField = new ScriptField(entry.getKey(), entry.getValue());
                query.addScriptField(scriptField);
            }
        }

        // 查询数据
        SearchHits<T> search = ES.search(query, clazz);
        if (isSortLoca) {// 是否按坐标排序
            for (SearchHit<T> searchHit : search) {
                T content = searchHit.getContent();
                list.add(content);

                // 添加距离信息
                List<Object> sortValues = searchHit.getSortValues();
                if (sortValues != null && sortValues.size() > 0) {
                    EsBaseBean base = (EsBaseBean)content;
                    base.setDist(Double.parseDouble(sortValues.get(0).toString()));
                }
            }
        } else {
            for (SearchHit<T> searchHit : search) {
                list.add(searchHit.getContent());
            }
        }

        // 设置数据总数
        long total = search.getTotalHits();
        list.setTotal(total);
        // 设置分页
        if (isPage) {
            list.setPageNum(pageNum);
            list.setPageSize(pageSize);
            long pages = total / pageSize;
            if (pages * pageSize < total) {
                pages++;
            }
            list.setPages((int)pages);
        }
        return list;
    }

    /**
     * 统计数据
     * 
     * @param clazz     索引类
     * @return
     * @throws Exception
     */
    public static long count(Class<?> clazz) throws Exception {
        return count(clazz, QueryBuilders.boolQuery());
    }

    /**
     * 统计数据
     * 
     * @param clazz     索引类
     * @param builder   查询条件{@link com.sojson.util.search.es.util.EsSearchUtil#build()}
     * @return
     * @throws Exception
     */
    public static long count(Class<?> clazz, AbstractQueryBuilder<?> builder) throws Exception {
        return count(clazz, builder, null, null).getHits().getTotalHits().value;
    }

    /**
     * 统计数据
     * 
     * @param clazz     索引类
     * @param aggrs     聚合函数{@link org.elasticsearch.search.aggregations.AggregationBuilders#max("别名")#field("字段名")}
     * @return
     * @throws Exception
     */
    public static SearchResponse count(Class<?> clazz, List<AggregationBuilder> aggrs) throws Exception {
        return count(clazz, QueryBuilders.boolQuery(), aggrs);
    }

    /**
     * 统计数据
     * 
     * @param clazz     索引类
     * @param scripts   脚本字段(key为返回的字段名)(由其他字段计算得到){@link com.sojson.util.search.es.util.EsSearchUtil#createScript()}
     * @return
     * @throws Exception
     */
    public static SearchResponse count(Class<?> clazz, Map<String, Script> scripts) throws Exception {
        return count(clazz, QueryBuilders.boolQuery(), scripts);
    }

    /**
     * 统计数据
     * 
     * @param clazz     索引类
     * @param builder   查询条件{@link com.sojson.util.search.es.util.EsSearchUtil#build()}
     * @param aggrs     聚合函数{@link org.elasticsearch.search.aggregations.AggregationBuilders#max("别名")#field("字段名")}
     * @return
     * @throws Exception
     */
    public static SearchResponse count(Class<?> clazz, AbstractQueryBuilder<?> builder, List<AggregationBuilder> aggrs)
        throws Exception {
        return count(clazz, builder, aggrs, null);
    }

    /**
     * 统计数据
     * 
     * @param clazz     索引类
     * @param builder   查询条件{@link com.sojson.util.search.es.util.EsSearchUtil#build()}
     * @param scripts   脚本字段(key为返回的字段名)(由其他字段计算得到){@link com.sojson.util.search.es.util.EsSearchUtil#createScript()}
     * @return
     * @throws Exception
     */
    public static SearchResponse count(Class<?> clazz, AbstractQueryBuilder<?> builder, Map<String, Script> scripts)
        throws Exception {
        return count(clazz, builder, null, scripts);
    }

    /**
     * 统计数据
     * 
     * @param clazz     索引类
     * @param aggrs     聚合函数{@link org.elasticsearch.search.aggregations.AggregationBuilders#max("别名")#field("字段名")}
     * @param scripts   脚本字段(key为返回的字段名)(由其他字段计算得到){@link com.sojson.util.search.es.util.EsSearchUtil#createScript()}
     * @return
     * @throws Exception
     */
    public static SearchResponse count(Class<?> clazz, List<AggregationBuilder> aggrs, Map<String, Script> scripts)
        throws Exception {
        return count(clazz, QueryBuilders.boolQuery(), aggrs, scripts);
    }

    /**
     * 统计数据
     * 
     * @param clazz     索引类
     * @param builder   查询条件{@link com.sojson.util.search.es.util.EsSearchUtil#build()}
     * @param aggrs     聚合函数{@link org.elasticsearch.search.aggregations.AggregationBuilders#max("别名")#field("字段名")}
     * @param scripts   脚本字段(key为返回的字段名)(由其他字段计算得到){@link com.sojson.util.search.es.util.EsSearchUtil#createScript()}
     * @return
     * @throws Exception
     */
    public static SearchResponse count(Class<?> clazz, AbstractQueryBuilder<?> builder, List<AggregationBuilder> aggrs,
        Map<String, Script> scripts) throws Exception {
        // 包装build
        builder = packageBuild(builder);

        // 获取索引名称
        String indexName = ES.getIndexCoordinatesFor(clazz).getIndexName();

        // 创建请求
        SearchRequest request = new SearchRequest(indexName);

        // filter不会对条件打分,效率更高
        if (log.isDebugEnabled()) {
            log.debug("{{}:{}}", "\"query\"", Objects.requireNonNull(builder));
        }

        // 设置查询条件
        SearchSourceBuilder source = new SearchSourceBuilder().query(builder);
        // 统计不需要数据,设置size=0
        source = source.size(0);
        // 设置允许超过1万条查询限制
        source = source.trackTotalHits(true);
        // 设置sum聚合
        if (aggrs != null && aggrs.size() > 0) {
            for (AggregationBuilder aggr : aggrs) {
                source = source.aggregation(aggr);
            }
        }
        // 设置脚本字段
        if (scripts != null && scripts.size() > 0) {
            Set<Entry<String, Script>> entrySet = scripts.entrySet();
            for (Entry<String, Script> entry : entrySet) {
                source.scriptField(entry.getKey(), entry.getValue());
            }
        }
        request.source(source);

        // 查询数据
        SearchResponse search = CLIENT.search(request, RequestOptions.DEFAULT);

        // 返回数据
        return search;
    }

    /**
     * 创建查询条件
     * 
     * @return
     * @throws Exception
     */
    public static EsSearchUtil createQuery() throws Exception {
        return new EsSearchUtil(QueryBuilders.boolQuery());
    }

    /**
     * 复制排序规则
     * 
     * @param sorts
     * @return
     */
    private static List<SortBuilder<?>> copySort(List<OrderItem> sorts) {
        if (sorts == null || sorts.size() < 1) {
            return new ArrayList<>(0);
        }

        List<SortBuilder<?>> builders = new ArrayList<>(sorts.size());
        for (OrderItem sort : sorts) {
            FieldSortBuilder order =
                SortBuilders.fieldSort(sort.getColumn()).order(sort.isAsc() ? SortOrder.ASC : SortOrder.DESC);
            builders.add(order);
        }
        return builders;
    }

    /**
     * 包装Builder
     * 
     * @param builder
     * @return
     */
    private static BoolQueryBuilder packageBuild(AbstractQueryBuilder<?> builder) {
        // filter不会对条件打分,效率更高
        return QueryBuilders.boolQuery().filter(builder);
    }

    /**
     * 获取索引名称
     * 
     * @param name  索引名称
     * @return
     */
    public static String getIndexName(String name) {
        return StringUtil.toUnderScoreCase(StringUtil.getKeyBase(name));
    }

}